# Autor chenfeng
#!/usr/bin/env Python
# coding=utf-8

#relation
import json
from paddlenlp import Taskflow

def openreadtxt(file_name):
    data = []
    file = open(file_name,'r',encoding='UTF-8')  #打开文件
    file_data = file.readlines() #读取所有行
    for row in file_data:
        data.append(row) #将每行数据插入data中
    return data

data_input=openreadtxt('test.txt')



# schema = {"公司":"高管"} # 该schema下，‘黄峥，拼多多’这条语句识别是空
schema = {"高管":"公司"}  # 该schema下，‘哔哩哔哩公司’和‘黄峥，拼多多’ 2条语句都可正确识别
# 微调后的模型
few_ie = Taskflow('information_extraction', schema=schema, batch_size=8, task_path='../checkpoint_gx/model_best')
# 默认模型
# few_ie = Taskflow('information_extraction', schema=schema, batch_size=8)

results=few_ie(data_input)

with open("result.txt", "w+",encoding='UTF-8') as f:    #a :   写入文件，若文件不存在则会先创建再写入，但不会覆盖原文件，而是追加在文件末尾
    for result in results:
        line = json.dumps(result, ensure_ascii=False)  #对中文默认使用的ascii编码.想输出真正的中文需要指定ensure_ascii=False
        f.write(line + "\n")

print("数据结果已导出")
for idx, text in enumerate(data_input):
    print(data_input[idx],results[idx])